{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from helpers import make_state_trace, identify_system, random_input, plot_eigenvalues\n",
    "from mpl_toolkits.axes_grid1.inset_locator import inset_axes\n",
    "from sklearn.linear_model import OrthogonalMatchingPursuit\n",
    "%matplotlib notebook\n",
    "\n",
    "# Code written by Kuan-Yun Lee"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Outlier Removal via OMP\n",
    "\n",
    "This notebook will guide you through the process of removing outliers for problems of the form\n",
    "$$ y \\approx Xw. $$\n",
    "\n",
    "To begin, we generate data for our problem. All the values in our code can be changed, so you can play around with the numbers if you want. We model $w$ as a random vector consisting of uniform random variables in $[2,20]$, and model the data points $X$ as a matrix with uniform entries in $[0,10]$. We observe corrupted versions of the true observations $y = Xw + Z$, where $Z$ is a multivariate Gaussian random variable with diagonal covariance depending on whether the observation is an outlier. In particular, we choose uniformly $num\\_outliers$ amount of outliers, and for these outliers, the variance is set to $100$, while for other points, the variance is set to $0.1$. In order to make the outliers exaggerated for us to see the effect of our algorithm, we add $40$ to outliers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --------------------------------------------------------------------\n",
    "\n",
    "def generate_data(n, d, num_outliers):\n",
    "    \"\"\"\n",
    "    Generates the data for our problem.\n",
    "    \n",
    "    args:\n",
    "      n: number of samples\n",
    "      d: dimension of samples\n",
    "      num_outliers: how many outliers are in the noisy observations;\n",
    "        we model outliers by adding large Gaussian noise\n",
    "                      \n",
    "    returns:\n",
    "      data: the nxd data matrix X\n",
    "      w_star: the underlying dx1 matrix W\n",
    "      observations: nx1 vector of noisy observations with outliers\n",
    "      \n",
    "    \"\"\"\n",
    "    \n",
    "    w_star = np.random.uniform(2.0,20.0,(d,1))\n",
    "    data = np.random.uniform(0.0,10.0,(n,d))\n",
    "    y_true = np.dot(data,w_star)\n",
    "    ind = np.random.choice(np.arange(n), num_outliers, False)\n",
    "    vec = 0.1*np.ones((n,1))\n",
    "    for a in ind:\n",
    "        vec[a] = 100.000\n",
    "    observations = np.random.multivariate_normal(\n",
    "                        np.ndarray.flatten(y_true), np.diagflat(vec))\n",
    "    for a in ind:\n",
    "        observations[a] += 40.000\n",
    "    return data, w_star, observations\n",
    "\n",
    "# --------------------------------------------------------------------"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Greedy OMP\n",
    "\n",
    "We implement our OMP with scikit-learn's OrthogonalMatchingPursuit function.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --------------------------------------------------------------------\n",
    "\n",
    "def greedy_OMP(X_aug, y, k, w_star):\n",
    "    \"\"\"\n",
    "    Performs orthogonal matching pursuit.\n",
    "    \n",
    "    args:\n",
    "      X_aug: augmented data matrix, i.e., X_aug = [X I]\n",
    "      y: noisy observations with outliers\n",
    "      k: how many non-zero entries in solution of OMP; is equal to\n",
    "        the number of iterations in the OMP algorithm\n",
    "                      \n",
    "    returns:\n",
    "      idx: the size k set of ignored data points\n",
    "      plot_points: the prediction points for the two points x=0 and \n",
    "        x = 10.0; for use of plotting\n",
    "       remaining_residual: amount of residual remaining\n",
    "        \n",
    "    \"\"\"\n",
    "    \n",
    "    reg = OrthogonalMatchingPursuit(n_nonzero_coefs=k, fit_intercept=False).fit(X_aug,y)\n",
    "    idx = [ j for (i,j) in zip(reg.coef_,range(1,len(reg.coef_)+1)) \n",
    "                       if (i>0.0000001 or i<-0.0000001) ]\n",
    "    \n",
    "    test_mat = 0.0*np.zeros((2,X_aug.shape[1]))\n",
    "    test_mat[1,0] = 10.0\n",
    "    plot_points = reg.predict(test_mat)\n",
    "    \n",
    "    y = np.reshape(y,(y.shape[0],1))\n",
    "    err_vec = (y - np.reshape(reg.predict(X_aug),(y.shape[0],1)))**2\n",
    "    remaining_residual = np.sum(err_vec)\n",
    "    \n",
    "    estimate_residual = (np.reshape(reg.predict(X_aug),(y.shape[0],1)) \n",
    "                                - np.reshape(X_aug[:,0]*w_star, (y.shape[0],1)))**2\n",
    "    \n",
    "    #Comment below to plot two lines, and uncomment the next line \n",
    "    return idx, plot_points, remaining_residual, np.zeros(np.asarray(err_vec).shape), np.log(np.asarray(estimate_residual))\n",
    "    #return idx, plot_points, remaining_residual, np.log(np.asarray(err_vec)), np.log(np.asarray(estimate_residual))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plotting the graph\n",
    "\n",
    "For each iteration of the algorithm, we plot the changes in our remaining points, the estimated line, and the remaining residue."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_func(X_aug, X_true, y, idxes, errs, est_plot, err_vec=[], est_err_vec=[]):\n",
    "    \"\"\"\n",
    "    Plots functions.\n",
    "    \n",
    "    args:\n",
    "      X_aug: augmented data matrix, i.e., X_aug = [X_true I]\n",
    "      X_true: the true data matrix\n",
    "      y: noisy observations with outliers\n",
    "      idxes: columns chosen by OMP\n",
    "      errs: errs from OMP\n",
    "      est_plot: estimated values at x=0 and x=10 for plotting of line\n",
    "      err_vec: vector of errors with respect to true fit\n",
    "      est_err_vec: vector of errors with respect to estimated fit\n",
    "        \n",
    "    \"\"\"\n",
    "    X_copy = X_true.copy()\n",
    "    y_copy = y.copy()\n",
    "    y_copy2 = y.copy()\n",
    "    X_deleted = []\n",
    "    Y_deleted = []\n",
    "    \n",
    "    err_vec_copy = err_vec.copy()\n",
    "    est_err_vec_copy = est_err_vec.copy()\n",
    "    err_vec_deleted = [None]*len(err_vec)\n",
    "    est_err_vec_deleted = [None]*len(est_err_vec)\n",
    "    \n",
    "    \n",
    "    for a in idxes: # Seperate the removed outliers from the data\n",
    "        if a!=1:\n",
    "            X_deleted.append(X_copy[a-2,0])\n",
    "            Y_deleted.append(y_copy2[a-2])\n",
    "            X_copy[a-2] = None\n",
    "            y_copy[a-2] = None \n",
    "            \n",
    "            err_vec_deleted[a-2] = err_vec_copy[a-2]\n",
    "            est_err_vec_deleted[a-2] = est_err_vec_copy[a-2]\n",
    "            err_vec[a-2] = None\n",
    "            est_err_vec[a-2] = None\n",
    "            \n",
    "    \n",
    "    if d==1: # Plot the linear fit and remaining points\n",
    "        fig, ax = plt.subplots(1,1,figsize = [10,10])\n",
    "        \n",
    "        plt.plot(X_copy[:,0],y_copy,'ro', label = 'Remaining points')\n",
    "        plt.plot(X_deleted, Y_deleted, 'gx', label = 'Removed points')\n",
    "        plt.plot(np.array([0,10.0]), est_plot, label = 'Linear fit')\n",
    "        plt.plot(np.array([0,10.0]), np.array([0,10.0*w_star]), label = 'True fit')\n",
    "        plt.title('Points and linear fit at iteration {}'.format(n-num+1))\n",
    "        plt.legend()\n",
    "        \n",
    "        axin = inset_axes(ax, width=\"30%\", height=\"30%\", loc=4, borderpad=2) \n",
    "        errs_vals = errs[:] + [None]*(num-1)\n",
    "        current_val = [None]*n\n",
    "        current_val[n-num] = errs_vals[n-num]\n",
    "        errs_vals[n-num] = None \n",
    "        \n",
    "        plt.title('Error over remaining points')\n",
    "        plt.xlim(0,n+1)\n",
    "        plt.ylim(-5,errs[0]+5)\n",
    "        plt.plot(np.arange(1,n+1,1), errs_vals, 'b.', label = 'Previous error')\n",
    "        plt.plot(np.arange(1,n+1,1), current_val, 'r.', label = 'Current error')\n",
    "        plt.plot()\n",
    "        plt.legend()\n",
    "        \n",
    "        \n",
    "        if len(err_vec)!=0 and len(est_err_vec)!=0:\n",
    "            axin = inset_axes(ax, width=\"70%\", height=\"10%\", loc=1, borderpad=2) \n",
    "            \n",
    "            axin.set_yticklabels([])\n",
    "            \n",
    "            plt.ylim(-0.5,0.5)\n",
    "            \n",
    "            \n",
    "            # UnComment the following lines to plot both\n",
    "            #plt.plot(err_vec, 0.2*np.ones((len(err_vec),)), 'r,')\n",
    "            #plt.plot(err_vec_deleted, 0.2*np.ones((len(err_vec),)), 'gx')\n",
    "            #plt.plot(est_err_vec, -0.2*np.ones((len(err_vec),)), 'r,')\n",
    "            #plt.plot(est_err_vec_deleted, -0.2*np.ones((len(err_vec),)), 'gx')\n",
    "            #plt.title(\"Top: Distances to Current Fit, Bottom: Distances to True Fit\")\n",
    "            #plt.xlabel(\"Log of distances\")\n",
    "            \n",
    "            #Use the following lines to plot just the distance to true fit, \n",
    "            #if you wish to plot both lines, comment the bottom and uncomment the above part\n",
    "            \n",
    "            plt.xlim(-5,10)\n",
    "            plt.plot(est_err_vec, np.zeros((len(err_vec),)), 'r,')\n",
    "            plt.plot(est_err_vec_deleted, np.zeros((len(err_vec),)), 'gx')\n",
    "            plt.title(\"Distances to True Fit\")\n",
    "            plt.xlabel(\"Log of distances\")\n",
    "        \n",
    "        plt.show()\n",
    "        \n",
    "        \n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Running the procedure\n",
    "\n",
    "First we generate our data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The plots will only work for dimension d = 1.\n",
    "n = 50\n",
    "d = 1\n",
    "num_outliers = 20\n",
    "\n",
    "X, w_star, y = generate_data(n,d,num_outliers) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we plot the initial points and the initial linear fit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "idxes = []\n",
    "errs = []\n",
    "X_true = X.copy()\n",
    "X_aug = np.concatenate((X,np.eye(X.shape[0])),axis =1)\n",
    "\n",
    "# plot the initial points and linear fit\n",
    "if d==1:\n",
    "    idxes, est_plot, _, _x, _y = greedy_OMP(X_aug,y, 1, w_star)\n",
    "    %matplotlib inline\n",
    "    fig, ax = plt.subplots(1,1,figsize = [10,10])\n",
    "    plt.plot(X,y,'ro')\n",
    "    plt.plot(np.array([0,10.0]), est_plot, label = 'Linear fit')\n",
    "    plt.plot(np.array([0,10.0]), np.array([0,10.0*w_star]), label = 'Optimal fit')\n",
    "    \n",
    "    plt.title('Initial points and linear fit')\n",
    "    plt.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we start the OMP procedure, and see what happens during each iteration. Note that here we let our procedure run until there are only two points left, or until the OMP algorithm cannot distinguish between the remaining points. This situation can happen when we have removed all the outliers and all the remaining points are on a line already."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num = n\n",
    "idxes_size = -1\n",
    "\n",
    "while(num>1): \n",
    "    idxes, est_plot, err, err_vec, est_err_vec = greedy_OMP(X_aug,y, n-num+2, w_star) \n",
    "    if idxes_size == len(idxes):\n",
    "        break\n",
    "    idxes_size = len(idxes)\n",
    "    errs.append(err)\n",
    "    \n",
    "    plot_func(X_aug, X_true, y, idxes, errs, est_plot, err_vec, est_err_vec)\n",
    "    \n",
    "    num = num - 1\n",
    "\n",
    "fig, ax = plt.subplots(1,1,figsize = [10,10])\n",
    "plt.plot(np.arange(1,len(errs)+1,1), errs,)\n",
    "plt.title(\"Error vs number of augmented bases\")\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example for system identification with outliers\n",
    "\n",
    "Although using ML libraries such as sklearn is helpful, in some cases we still need to do a bit more work as engineers to fix issues. In this example, we will see how outliers for systemID will affect our estimations, how our vanilla OMP will not always work, and how we can do some extra work to make it work.\n",
    "\n",
    "Now we consider Example 1 from the system ID Demo code. In particular, we'll use\n",
    "\\begin{align}\n",
    "A &= \\begin{bmatrix} 0 & 1 \\\\ 0.3 & 0.2 \\end{bmatrix}, & B &= \\begin{bmatrix} 0 \\\\ 1 \\end{bmatrix}.\n",
    "\\end{align}\n",
    "For more details, see Example 1 of the DemoSystemID code.\n",
    "\n",
    "Below is the exact same code from Example 1. As we can see, when there are no significant outliers, the algorithm there performs perfectly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_known = np.matrix([[0, 1],\n",
    "                     [0.3, 0.2]])\n",
    "b_known = np.matrix([[0],\n",
    "                     [1]])\n",
    "\n",
    "k = 100\n",
    "t = t = np.arange(k)\n",
    "\n",
    "u_trace = random_input(t, 1)\n",
    "\n",
    "x_trace = make_state_trace(a_known, b_known, u_trace)\n",
    "\n",
    "a_identified, b_identified = identify_system(x_trace, u_trace)\n",
    "print('identified A matrix:')\n",
    "print(a_identified)\n",
    "print('identified B matrix:')\n",
    "print(b_identified)\n",
    "print('Frobenius norm of matrix A and true A:',  \n",
    "      np.linalg.norm(a_identified-a_known, ord='fro', axis=None, keepdims=False))\n",
    "print('Frobenius norm of matrix B and true B:',  \n",
    "      np.linalg.norm(b_identified-b_known, ord='fro', axis=None, keepdims=False))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Effect of adding k outliers\n",
    "\n",
    "Now, we simply choose $k$ random points from $\\vec{x}$, and add $2000$ to both coordinates. Running the exact same code as above, there is already visible error in our estimates of $A$ and $b$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = 5\n",
    "\n",
    "corrupt_idx = np.random.choice(np.arange(0,x_trace.shape[0],1), size = (k,1), replace=False)\n",
    "x_trace_cor = x_trace.copy()\n",
    "for i in corrupt_idx:\n",
    "    x_trace_cor[i] = x_trace_cor[i]+2000\n",
    "\n",
    "a_identified_cor, b_identified_cor = identify_system(np.asarray(x_trace_cor), u_trace)\n",
    "print('identified A matrix:')\n",
    "print(a_identified_cor)\n",
    "print('identified B matrix:')\n",
    "print(b_identified_cor)\n",
    "print('Frobenius norm of matrix A and true A:',  \n",
    "      np.linalg.norm(a_identified_cor-a_known, ord='fro', axis=None, keepdims=False))\n",
    "print('Frobenius norm of matrix B and true B:',  \n",
    "      np.linalg.norm(b_identified_cor-b_known, ord='fro', axis=None, keepdims=False))\n",
    "plot_eigenvalues([a_known, a_identified_cor], ['true eigenvalues', 'estimated eigenvalues'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Effect of running OMP to remove the outlier\n",
    "\n",
    "Next, let's use our OMP algorithm to try to remove the outlier. As we can see, this doesn't work in this case (you can play with some examples to see when the algorithm still works). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_aug = np.concatenate((x_trace_cor[0:u_trace.shape[0],:], u_trace), axis = 1)\n",
    "X_aug = np.concatenate((X_aug, np.eye(u_trace.shape[0])), axis = 1)\n",
    "y = x_trace_cor[1:u_trace.shape[0]+1,:]\n",
    "\n",
    "\n",
    "\n",
    "reg1 = OrthogonalMatchingPursuit(n_nonzero_coefs=6, fit_intercept=False).fit(X_aug,y[:,0])\n",
    "reg2 = OrthogonalMatchingPursuit(n_nonzero_coefs=6, fit_intercept=False).fit(X_aug,y[:,1])\n",
    "X_0 = np.concatenate((np.eye(3), np.zeros((3,X_aug.shape[0])) ), axis = 1)\n",
    "a1 = np.reshape(reg1.predict(X_0),(1,3))\n",
    "a2 = np.reshape(reg2.predict(X_0),(1,3))\n",
    "a_mat = np.concatenate((a1, a2), axis = 0)\n",
    "\n",
    "\n",
    "idx1 = [ j for (i,j) in zip(reg1.coef_,range(1,len(reg1.coef_)+1)) \n",
    "                       if (i>0.0000001 or i<-0.0000001) ]\n",
    "idx2 = [ j for (i,j) in zip(reg2.coef_,range(1,len(reg2.coef_)+1)) \n",
    "                       if (i>0.0000001 or i<-0.0000001) ]\n",
    "\n",
    "new_identified_a = np.transpose(a_mat[0:2,0:2])\n",
    "new_identified_b = np.reshape(a_mat[:,2],(2,1))\n",
    "print('identified A matrix:')\n",
    "print(new_identified_a)\n",
    "print('identified B matrix:')\n",
    "print(new_identified_b)\n",
    "print('Frobenius norm of matrix A and true A:',  \n",
    "      np.linalg.norm(new_identified_a-a_known, axis=None, keepdims=False))\n",
    "print('Frobenius norm of matrix B and true B:',  \n",
    "      np.linalg.norm(new_identified_b-b_known, axis=None, keepdims=False))\n",
    "plot_eigenvalues([a_known, new_identified_a], ['true eigenvalues', 'OMP-estimated eigenvalues'])\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we see that removing the outlier with OMP doesn't work as amazingly as we hoped.\n",
    "\n",
    "## Fixing the issue\n",
    "To solve this issue, we implement a neighboring-OMP algorithm. The reason is that for this problem of system identification with state observations, we don't only have outliers in the observations or target vector --- they also exist as part of the data or features. To overcome this, we simply modify the algorithm such that when a basis $e_j$ is chosen, we also choose the next two bases to cover all the possible effects of the corrupted points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_idxes1 = []\n",
    "new_idxes2 = []\n",
    "\n",
    "for i in idx1:\n",
    "    if i>3 and i< len(x_trace_cor)-2:\n",
    "        new_idxes1.append(i-1)\n",
    "        new_idxes1.append(i)\n",
    "        new_idxes1.append(i+1)\n",
    "    elif i==len(x_trace_cor)-1:\n",
    "        new_idxes1.append(i-1)\n",
    "        new_idxes1.append(i)\n",
    "    elif i==len(x_trace_cor):\n",
    "        new_idxes1.append(i+1)\n",
    "        \n",
    "new_idxes1 = np.unique(np.asarray(new_idxes1))\n",
    "\n",
    "for i in idx2:\n",
    "    if i>3 and i< len(x_trace_cor)-2:\n",
    "        new_idxes2.append(i-1)\n",
    "        new_idxes2.append(i)\n",
    "        new_idxes2.append(i+1)\n",
    "    elif i==len(x_trace_cor)-1:\n",
    "        new_idxes2.append(i-1)\n",
    "        new_idxes2.append(i)\n",
    "    elif i==len(x_trace_cor):\n",
    "        new_idxes2.append(i-1)\n",
    "\n",
    "\n",
    "X1 = X_aug[new_idxes1, 0:3]\n",
    "X2 = X_aug[new_idxes2, 0:3]\n",
    "\n",
    "Ab1 = np.dot(np.dot(np.linalg.inv(np.dot(X1.T,X1)),X1.T),y[new_idxes1,0])\n",
    "Ab2 = np.dot(np.dot(np.linalg.inv(np.dot(X2.T,X2)),X2.T),y[new_idxes2,1])\n",
    "Ab = np.concatenate((Ab1,Ab2))\n",
    "estimated_A_ours = Ab[:,0:2]\n",
    "estimated_b_ours = Ab[:,2]\n",
    "\n",
    "print('identified A matrix:')\n",
    "print(estimated_A_ours)\n",
    "print('identified B matrix:')\n",
    "print(estimated_b_ours)\n",
    "print('Frobenius norm of matrix A and true A:',  \n",
    "      np.linalg.norm(estimated_A_ours-a_known, axis=None, keepdims=False))\n",
    "print('Frobenius norm of matrix B and true B:',  \n",
    "      np.linalg.norm(estimated_b_ours-b_known, axis=None, keepdims=False))\n",
    "plot_eigenvalues([a_known, estimated_A_ours], ['true eigenvalues', 'our estimated eigenvalues'])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we see that our method works nicely!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
